import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.metrics import (
    f1_score, accuracy_score, normalized_mutual_info_score,
    rand_score, adjusted_rand_score
)

# 1. 加载数据集
iris = load_iris()
X = iris.data  # 特征：花萼长度、花萼宽度、花瓣长度、花瓣宽度
y_true = iris.target  # 真实标签（用于评估，聚类本身无监督）
feature_names = iris.feature_names
target_names = iris.target_names

# 2. 实现 KMeans 类
class KMeans:
    def __init__(self, k, max_iter=100, tol=1e-4):
        """
        k: 聚类数
        max_iter: 最大迭代次数
        tol: 中心点变化阈值（提前停止依据）
        """
        self.k = k
        self.max_iter = max_iter
        self.tol = tol
        self.centroids = None  # 存储聚类中心
        self.labels = None     # 存储样本聚类标签

    def _euclidean_distance(self, x1, x2):
        """计算欧式距离"""
        return np.sqrt(np.sum((x1 - x2) ** 2, axis=1))

    def fit(self, X):
        # 随机初始化聚类中心（从样本中选 k 个）
        init_indices = np.random.choice(len(X), self.k, replace=False)
        self.centroids = X[init_indices]

        for _ in range(self.max_iter):
            # 1. 分配样本到最近的聚类中心
            distances = np.array([self._euclidean_distance(x, self.centroids) for x in X])
            self.labels = np.argmin(distances, axis=1)

            # 2. 更新聚类中心
            new_centroids = np.array([X[self.labels == i].mean(axis=0) for i in range(self.k)])

            # 3. 检查是否提前停止（中心变化小于 tol）
            if np.sum(np.abs(new_centroids - self.centroids)) < self.tol:
                break
            self.centroids = new_centroids

# 3. 可视化聚类结果（取前 2 个特征，方便 2D 展示）
def visualize_clusters(X, y_pred, y_true, feature_names, title):
    plt.figure(figsize=(8, 5))
    # 聚类结果
    plt.scatter(X[:, 0], X[:, 1], c=y_pred, cmap='viridis', edgecolors='k', label='Predicted Clusters')
    # 真实标签（辅助对比）
    plt.scatter(X[:, 0], X[:, 1], c=y_true, cmap='rainbow', alpha=0.3, s=100, label='True Labels')
    plt.xlabel(feature_names[0])
    plt.ylabel(feature_names[1])
    plt.title(title)
    plt.legend()
    plt.show()

# 4. 计算评估指标（需注意：聚类无监督，真实标签仅用于对比评估）
def evaluate_clustering(y_true, y_pred):
    # F-measure（需指定平均策略，多分类常用'macro'）
    f_measure = f1_score(y_true, y_pred, average='macro')
    # 聚类版 ACC（需匹配标签顺序，和分类 ACC 逻辑不同，仅数值对比）
    acc = accuracy_score(y_true, y_pred)
    nmi = normalized_mutual_info_score(y_true, y_pred)
    ri = rand_score(y_true, y_pred)
    ari = adjusted_rand_score(y_true, y_pred)
    return {
        'F-measure': f_measure,
        'ACC': acc,
        'NMI': nmi,
        'RI': ri,
        'ARI': ari
    }

# 5. 主流程：训练 + 可视化 + 评估
if __name__ == "__main__":
    # 初始化模型（k=3 对应 iris 三类，可调整；max_iter 控制迭代，tol 控制提前停止）
    kmeans = KMeans(k=3, max_iter=50, tol=1e-3)
    kmeans.fit(X)
    y_pred = kmeans.labels

    # 可视化（用前两个特征）
    visualize_clusters(X[:, :2], y_pred, y_true, feature_names,
                      "KMeans Clustering on Iris Dataset (2D Projection)")

    # 评估指标
    metrics = evaluate_clustering(y_true, y_pred)
    print("聚类评估指标：")
    for metric, value in metrics.items():
        print(f"{metric}: {value:.4f}")